from typing import Any, Tuple, List, Optional, Sized

from scipy.ndimage.interpolation import zoom
import itertools
from torch.utils.data.sampler import Sampler
from torchvision import transforms
import torch
import numpy as np
from scipy import ndimage

import random
import cv2
from PIL import Image, ImageOps, ImageFilter
from scipy.ndimage import gaussian_filter
from torch import nn
from torch.nn import functional as F


def patients_to_slices(dataset, patiens_num):
    ref_dict = None
    if "ACDC" in dataset:
        ref_dict = {"3": 68, "7": 136,
                    "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312}
    elif "Prostate":
        ref_dict = {"2": 27, "4": 53, "8": 120,
                    "12": 179, "16": 256, "21": 312, "42": 623}
    else:
        print("Error")
    return ref_dict[str(patiens_num)]


class TwoStreamBatchSampler(Sampler):
    """Iterate two sets of indices

    An 'epoch' is one iteration through the primary indices.
    During the epoch, the secondary indices are iterated through
    as many times as needed.
    """

    def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size,
                 data_source: Optional[Sized]):
        super().__init__(data_source)
        self.primary_indices = primary_indices
        self.secondary_indices = secondary_indices
        self.secondary_batch_size = secondary_batch_size
        self.primary_batch_size = batch_size - secondary_batch_size

        assert len(self.primary_indices) >= self.primary_batch_size > 0
        assert len(self.secondary_indices) >= self.secondary_batch_size > 0

    def __iter__(self):
        primary_iter = iterate_once(self.primary_indices)
        secondary_iter = iterate_eternally(self.secondary_indices)
        return (
            primary_batch + secondary_batch
            for (primary_batch, secondary_batch) in zip(
            grouper(primary_iter, self.primary_batch_size),
            grouper(secondary_iter, self.secondary_batch_size),
        )
        )

    def __len__(self):
        return len(self.primary_indices) // self.primary_batch_size


def iterate_once(iterable):
    return np.random.permutation(iterable)


def iterate_eternally(indices):
    def infinite_shuffles():
        while True:
            yield np.random.permutation(indices)

    return itertools.chain.from_iterable(infinite_shuffles())


def grouper(iterable, n):
    "Collect data into fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3) --> ABC DEF"
    args = [iter(iterable)] * n
    return zip(*args)


def random_rot_flip(image, label):
    k = np.random.randint(0, 4)
    image = np.rot90(image, k)
    label = np.rot90(label, k)
    axis = np.random.randint(0, 2)
    image = np.flip(image, axis=axis).copy()
    label = np.flip(label, axis=axis).copy()
    return image, label


def random_rotate(image, label):
    angle = np.random.randint(-20, 20)
    image = ndimage.rotate(image, angle, order=0, reshape=False)
    label = ndimage.rotate(label, angle, order=0, reshape=False)
    return image, label


def color_jitter(image):
    if not torch.is_tensor(image):
        np_to_tensor = transforms.ToTensor()
        image = np_to_tensor(image)

    # s is the strength of color distortion.
    s = 1.0
    jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)
    return jitter(image)


def crop(img, mask, size: Tuple = (224, 224), ignore_value=255):
    w, h = img.size
    padw = size[0] - w if w < size[0] else 0
    padh = size[1] - h if h < size[1] else 0
    img = ImageOps.expand(img, border=(0, 0, padw, padh), fill=0)
    mask = ImageOps.expand(mask, border=(0, 0, padw, padh), fill=ignore_value)

    w, h = img.size
    x = random.randint(0, w - size[0])
    y = random.randint(0, h - size[1])
    img = img.crop((x, y, x + size[0], y + size[1]))
    mask = mask.crop((x, y, x + size[0], y + size[1]))

    return img, mask


def hflip(img, mask, p=0.5):
    if random.random() < p:
        img = img.transpose(Image.FLIP_LEFT_RIGHT)
        mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
    return img, mask


def normalize(img, mask=None):
    img = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])(img)
    if mask is not None:
        mask = torch.from_numpy(np.array(mask)).long()
        return img, mask
    return img


def normalize_two(img, mask=None):
    img = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])(img)
    if mask is not None:
        mask = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])(mask)
        return img, mask
    return img


def de_normalizer(image):
    MEAN = [-mean / std for mean, std in zip([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]
    STD = [1 / std for std in [0.229, 0.224, 0.225]]
    denormalizer = transforms.Normalize(mean=MEAN, std=STD)
    image = denormalizer(image)
    image = image.mul(255.0).byte()
    image = image.numpy().transpose((1, 2, 0))
    return image


def rand_resize(img, mask, ratio_range):
    w, h = img.size
    long_side = random.randint(int(max(h, w) * ratio_range[0]), int(max(h, w) * ratio_range[1]))

    if h > w:
        oh = long_side
        ow = int(1.0 * w * long_side / h + 0.5)
    else:
        ow = long_side
        oh = int(1.0 * h * long_side / w + 0.5)

    img = img.resize((ow, oh), Image.BILINEAR)
    mask = mask.resize((ow, oh), Image.NEAREST)
    return img, mask


def resize(img, mask, size: Tuple):
    img = img.resize(size, Image.BILINEAR)
    mask = mask.resize(size, Image.NEAREST)
    return img, mask


def resize_two(img, mask, size: Tuple):
    img = img.resize(size, Image.BILINEAR)
    mask = mask.resize(size, Image.BILINEAR)
    return img, mask


def rand_resize_two(img, mask, ratio_range):
    w, h = img.size
    long_side = random.randint(int(max(h, w) * ratio_range[0]), int(max(h, w) * ratio_range[1]))

    if h > w:
        oh = long_side
        ow = int(1.0 * w * long_side / h + 0.5)
    else:
        ow = long_side
        oh = int(1.0 * h * long_side / w + 0.5)

    img = img.resize((ow, oh), Image.BILINEAR)
    mask = mask.resize((ow, oh), Image.BILINEAR)
    return img, mask


def blur(img, p=0.5):
    if random.random() < p:
        sigma = np.random.uniform(0.1, 2.0)
        img = img.filter(ImageFilter.GaussianBlur(radius=sigma))
    return img


class Transform_Gray(object):
    # 黑白图像的增广
    def __init__(self, output_size: Tuple = (224, 224)):
        self.output_size = output_size

    def __call__(self, image: np.ndarray, mask: np.ndarray):

        if random.random() > 0.5:
            image, mask = random_rot_flip(image, mask)
        elif random.random() > 0.5:
            image, mask = random_rotate(image, mask)
        x, y = image.shape

        image = zoom(image, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        mask = zoom(mask, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        image = torch.from_numpy(image.astype(np.float32)).unsqueeze(0)
        mask = torch.from_numpy(mask.astype(np.uint8))

        sample = {'image': image, 'mask': mask}
        return sample


class WeakStrongAugment_Gray(object):
    """returns weakly and strongly augmented images

    Args:
        object (tuple): output size of network
    """

    def __init__(self, output_size: Tuple = (224, 224)):
        self.output_size = output_size

    def __call__(self, image: np.ndarray, mask: np.ndarray):
        image = self.resize(image)
        mask = self.resize(mask)
        # weak augmentation is rotation / flip
        image_weak, mask = random_rot_flip(image, mask)
        # strong augmentation is color jitter
        image_strong = color_jitter(image_weak).type("torch.FloatTensor")
        # fix dimensions
        image = torch.from_numpy(image.astype(np.float32)).unsqueeze(0)
        image_weak = torch.from_numpy(image_weak.astype(np.float32)).unsqueeze(0)
        mask = torch.from_numpy(mask.astype(np.uint8))

        sample = {
            "image": image,
            "image_weak": image_weak,
            "image_strong": image_strong,
            "label_aug": mask,
        }
        return sample

    def resize(self, image):
        x, y = image.shape
        return zoom(image, (self.output_size[0] / x, self.output_size[1] / y), order=0)


class Transform_RGB():
    # 彩色图像增广

    def __init__(self, mode="train", out_size=(224, 224)) -> None:
        self.mode = mode
        self.out_size = out_size

    def __call__(self, image: Image.Image, mask: Image.Image) -> Any:
        """
        输入彩色图片
        :param image:
        :param mask:
        :return:
        """
        if self.mode == 'test':
            image, mask = resize(image, mask, self.out_size)
            image, mask = normalize(image, mask)
            return {
                "image": image,
                "mask": mask
            }

        image, mask = rand_resize(image, mask, (0.5, 2.0))
        image, mask = crop(image, mask, self.out_size, 255)
        image, mask = hflip(image, mask, p=0.5)

        image, mask = normalize(image, mask)
        return {
            "image": image,
            "mask": mask
        }


class Transform_Two_RGB():
    # 彩色图像增广

    def __init__(self, mode="train", out_size=(224, 224)) -> None:
        self.mode = mode
        self.out_size = out_size

    def __call__(self, image: Image.Image, mask: Image.Image) -> Any:
        """
        输入彩色图片
        :param image:
        :param mask:
        :return:
        """
        if self.mode == 'test':
            image, mask = resize_two(image, mask, self.out_size)
            image, mask = normalize_two(image, mask)
            return {
                "image": image,
                "mask": mask
            }

        image, mask = rand_resize_two(image, mask, (0.5, 2.0))
        image, mask = crop(image, mask, self.out_size, 255)
        image, mask = hflip(image, mask, p=0.5)

        image, mask = normalize_two(image, mask)
        return {
            "image": image,
            "mask": mask
        }
